from NNets import *
from BaseNet import *
from deep_learning import *
from cv_helper import imread_cv, show_plt
from os.path import dirname
from torch.nn import CrossEntropyLoss


def p(*lst, **dic) -> None:
    """快速打印调试"""
    for i in lst:
        print(type(i), i, sep='\n', end="\n\n")
    for k, v in dic.items():
        print(type(k), k, type(v), v, sep='\n', end="\n\n")


# 神经网络文件保存位置
FILE_NAME: str = dirname(__file__) + "\\LeNet"
# 要识别的图像的位置
IMG_NAME: str = dirname(__file__) + "\\tes\\61.png"
# 是否做训练
TRAIN: bool = True
# 超参数
HYPER_PARAMETER: dict = {
    "BATCH_SIZE": 200,
    "EPOCH": 3,
    "LEARNING_RATE": 0.001,
    "LOSS_FUNC": CrossEntropyLoss(),
    "OPTIMIZER": "sgd",
}

if __name__ == "__main__":
    mod = loadModel(FILE_NAME, LeNet, HYPER_PARAMETER)
    myTest(mod)
    if TRAIN:
        myTrains(mod)
        myTest(mod)
        mod.saveModel(FILE_NAME)

    img = imread_cv(IMG_NAME)
    # if UNLOCK_DEVICE:
    #    img = img.to(DEVICE)
    outputs = distinguish_withMark(mod, img)
    show_plt(outputs)

    # compareMod(mod)
